{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "给定对称加密算法Enc，Dec，两个阶为$p$的群$G_1$和$G_2$，双线性映射 $ e: G_1 \\times G_1 \\rightarrow G_T $ ，$g$为群$G_1$的生成元，以及两个哈希函数$h_1, h_2$其中$h_1$可以将输入值映射到$G_1$，此外，假设所有的用户属于一个集合，并存在一个管理员UM，具体的步骤如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Crypto.Cipher import AES\n",
    "from Crypto.Random import get_random_bytes\n",
    "from charm.toolbox.pairinggroup import PairingGroup, G1, ZR, pair as e\n",
    "from Util import read_word_as_bytes, obj2bytes\n",
    "from hashlib import sha256, md5\n",
    "\n",
    "group = PairingGroup('SS512')\n",
    "g = group.random(G1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enc和Dec定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Enc(key, bdata):\n",
    "    if len(key) % 16 != 0:\n",
    "        raise ValueError('the bytes of key must be multiples of 16')\n",
    "    if len(bdata) % 16 != 0:\n",
    "        raise ValueError('the bytes of data must be multiples of 16')\n",
    "    cipher = AES.new(key, AES.MODE_ECB)\n",
    "    return cipher.encrypt(bdata)\n",
    "\n",
    "def Dec(key, bdata):\n",
    "    if len(key) % 16 != 0:\n",
    "        raise ValueError('the bytes of key must be multiples of 16')\n",
    "    if len(bdata) % 16 != 0:\n",
    "        raise ValueError('the bytes of data must be multiples of 16')\n",
    "    cipher = AES.new(key, AES.MODE_ECB)\n",
    "    return cipher.decrypt(bdata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enc 和 Dec测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def EncDecTest():\n",
    "    key = get_random_bytes(16)\n",
    "    # data = 'hello' # asser error\n",
    "    data = get_random_bytes(16)\n",
    "    secret = Enc(key, data)\n",
    "    plaintext = Dec(key, secret)\n",
    "    assert data == plaintext\n",
    "\n",
    "EncDecTest() # no error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$h_1: \\rightarrow G_1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def h1(b):\n",
    "    return group.hash(sha256(b).hexdigest(), G1)\n",
    "\n",
    "def h2(b):\n",
    "    return md5(b).hexdigest()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setup \\\n",
    "管理员UM选取随机数$x \\in Z_p^*$并将其作为自己的私钥，$k_{UM} = x$。对于某个用户U，UM首先选择$ x \\in Z_p^* $并计算 $ ComK_U = g^{x / x_U} \\in G_1$，然后将$x_U$秘密发送给用户U，同时将$ComK_U$发送给服务器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup():\n",
    "    x, x_U = group.random(ZR), group.random(ZR)\n",
    "    return {'k_UM': x, 'U': {'x_U': x_U, 'ComK_U': g ** (x / x_U)}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "build_index \\\n",
    "当用户U希望为关键词$w$生成索引时，首先选取随机数$r \\in Z_p^*$，然后将$ h_1\\left( w \\right) ^r $发送给服务器，随后服务期根据用户U对应的$ComK_U$计算$ e'=e\\left( h_1\\left( w \\right) ^r,ComK_U \\right)  $并将计算结果返回给用户。最后用户U计算$ k=h_2\\left( e'^{x_U/r} \\right)  $，得到关键词对应的索引$ I=\\left( m,Enc_k\\left( m \\right) \\right)  $，其中m为随机消息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_index(w, U):\n",
    "    r = group.random(ZR)\n",
    "    e_minute = e(h1(w) ** r, U['ComK_U'])\n",
    "    k = h2(obj2bytes(e_minute ** (U['x_U'] / r))).encode()\n",
    "    m = get_random_bytes(16)\n",
    "    return m, Enc(k, m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "generate_trapdoor \\\n",
    "当用户U想要检索关键词$w$时，为其计算陷门$ q=h_1\\left( w \\right) ^{x_U} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_trapdoor(w, U):\n",
    "    return h1(w) ** U['x_U']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "search \\\n",
    "根据用户U对应的$ComK_U$服务器首先计算$ k'=h_2\\left( e\\left( q,ComK_U \\right) \\right)  $，如果对于某条索引$ I=\\left( A,B \\right)  $，满足$ A=Dec_{k'}\\left( B \\right)  $，则该索引对应的关键词即为检索关键词。\\\n",
    "正确性证明：\\\n",
    "根据用户U的检索陷门，服务器可以得到一个解密秘钥：\n",
    "$$\n",
    "k'=h_2\\left( e\\left( q,ComK_U \\right) \\right) =h_2\\left( e\\left( h_1\\left( w \\right) ^{x_U},g^{x/x_U} \\right) \\right) =h_2\\left( e\\left( h_1\\left( w \\right) ,g \\right) ^x \\right)\n",
    "$$\n",
    "而对于用户Y上传的关键词$w$对应的索引，其加密秘钥为：\n",
    "$$\n",
    "k=h_2\\left( e'^{x_Y/r} \\right) =h_2\\left( e\\left( h_1\\left( w \\right) ^r,g^{x/x_Y} \\right) ^{x_Y/r} \\right) =h_2\\left( e\\left( h_1\\left( w \\right) ,g \\right) ^x \\right)\n",
    "$$\n",
    "可见，计算过程中用户的私钥被抵消了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(U, q, W_I):\n",
    "    k_minute = h2(obj2bytes(e(q, U['ComK_U']))).encode()\n",
    "    A, B = W_I\n",
    "    return A == Dec(k_minute, B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "params:  {'k_UM': 126059215243568542840966094678466802177637048183, 'U': {'x_U': 474388946183920585838392423799708302882447719683, 'ComK_U': [7006855254724661247936686240329011000912552848184446948688537247756642108650410724970843013665021090259293586647357587017725823532468618025001629158770699, 1379018327658185249172697926847885055542218549263392966920065315285842231934128412448863932349751865892387925656735322092299456352438844073305613579029319]}}\n",
      "Search happiness, result: Yes\n"
     ]
    }
   ],
   "source": [
    "wl = read_word_as_bytes('../test_files/happiness.txt')\n",
    "kw = b'happiness'\n",
    "params = setup()\n",
    "print('params: ', params)\n",
    "Idx = [build_index(w, params['U']) for w in wl]\n",
    "q = generate_trapdoor(kw, params['U'])\n",
    "print(f'Search {kw.decode()}, result: ', end = '')\n",
    "for W_I in Idx:\n",
    "    if search(params['U'], q, W_I):\n",
    "        print('Yes')\n",
    "        break\n",
    "else:\n",
    "    print('No')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e5030792b3492f6b12d94f1f48beca3d8e59ec05fd59d0aaaa48e684281ed297"
  },
  "kernelspec": {
   "display_name": "Python 3.7.9 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
